(defpackage :problem-generator
  (:use :cl)
  (:nicknames :pg)
  (:export :expand-and-fill
           :expand-and-fill-n-times
           :generator-error
           :generator-error-template
           :generator-error-extra
           :xbox-p))

(in-package :problem-generator)

(define-condition generator-error (error)
  ((template :initform "Unknown" :initarg :template :reader generator-error-template)
   (extra :initarg :extra :reader generator-error-extra))
  (:report (lambda (c s)
             (format s "Failed to generate a problem from template ~a: ~a"
                     (generator-error-template c) (generator-error-extra c)))))

(defstruct problem
  tree difficulty)
   
(defstruct constraint
  value override)

(defstruct (num-constraint (:include constraint) (:conc-name constraint-))
  (max 50)
  (negativep t)
  (fractionp nil))

(defstruct (op-constraint (:include constraint) (:conc-name constraint-))
  (ops '(+ - * /)))

(defun num ()
  (make-num-constraint))

(defun op ()
  (make-op-constraint))

(defun mul (&optional (divp t))
  (if divp
      (make-op-constraint '(* /))
      (make-op-constraint '(*))))

(defun add ()
  (make-op-constraint '(+ -)))

(defun gen-fractions (constraint)
  (declare (ignore constraint))
  (loop for i from 1 to 10
     nconc (loop for j from 1 to 10
              for f = (/ i j)
              when (typep f 'ratio)
              collect f)))

(defgeneric gen-possibles (thing))
(defgeneric pass-constraint (val thing)) 

(defmethod gen-possibles ((thing op-constraint))
  (u:randomize-list (constraint-ops thing)))

(defmethod gen-possibles ((thing num-constraint))
  (u:randomize-list
   (nconc
    (loop for i from 1 to (constraint-max thing)
       collect i
       if (constraint-negativep thing)
       collect (- i))
    (when (constraint-fractionp thing)
      (gen-fractions thing)))))

(defmethod pass-constraint (val (c num-constraint))
  (if (constraint-override c)
      (eql val (constraint-override c))
      (and (or (typep val 'ratio)
               (typep val 'integer))
           (not (zerop val))
           (or (constraint-negativep c)
               (plusp val))
           (or (and (constraint-fractionp c)
                    (typep val 'ratio)
                    (find val (gen-fractions c)))
               (typep val 'integer))
           (<= (- (constraint-max c)) val (constraint-max c)))))

(defmethod pass-constraint (val (c op-constraint))
  (if (constraint-override c)
      (eql val (constraint-override c))
      (find val (constraint-ops c))))

(defun pass-box-constraints (values constraints)
  (declare (ignore constraints))
  (and
   ;; dividing by itself is boring
   (or (not (eql (second values) '/))
       (/= (abs (first values)) (abs (third values))))
   ;; multiplying by 1 is boring.
   (or (not (find (second values) '(* /)))
       (and (/= (abs (first values)) 1)
            (/= (abs (third values)) 1)))))

(defgeneric walk-template (node &optional override-p))

(defun xbox-p (node)
  (and (typep node 'd:multi-argument-box)
       (eql (length (tree:children node)) 1)
       (typep (tree:child node) 'd:argument)
       (zerop (length (tree:children (tree:child node))))))

(defun dprint (thing)
;; (print thing)
  )

(defmacro with-values (children &body body)
  (u:with-gensyms (bk x i)
    `(let ((,bk (mapcar (lambda (,x) (constraint-value (d:quantity ,x))) ,children)))
       ,@body
       (loop for ,i in ,children
          for ,x in ,bk do
          (setf (constraint-value (d:quantity ,i)) ,x)))))

(defmethod walk-template ((node d:multi-argument-box) &optional override-p)
  (labels ((apply-value (val node)
;;              (assert (or (null (constraint-value (d:quantity node)))
;;                          (null val)
;;                          (eql val (constraint-value (d:quantity node))))
;;                      () "this is it ~a ~a" val (d:quantity node))
             (dprint `(ya ,val ,node))
             (setf (constraint-value (d:quantity node)) val))
           (apply-values (val parents)
             (dolist (p parents)
               (apply-value val p)))
           (calc (op n1 n2)
             (if (and (eq op '/) (zerop n2))
                 (error 'generator-error :template node :extra "Division by zero.")
                 (funcall op n1 n2)))
           (move-on ()
             (when (or (not (tree:parents node))
                       (every (lambda (p) (walk-template p override-p))
                              (mapcar 'tree:parent (tree:parents node))))
               ;; it's all in the side-effects
               (return-from walk-template t)))
           (skip-constraints-p (nodes)
             (and override-p
                  ;; if any children have overrides, then don't bother finding constraints.
                  (some (lambda (n)
                          (d:find-if-on-tree (lambda (n)
                                               (and (typep n 'd:argument)
                                                    (constraint-override (d:quantity n))))
                                             n))
                        nodes)))
           (valid-val (val nodes parents)
             (or (skip-constraints-p nodes)
                 (every (lambda (p) (pass-constraint val (d:quantity p)))
                        parents)))
           (valid-box (vals node)
             (or (skip-constraints-p (tree:children node))
                 (pass-box-constraints vals (mapcar 'd:quantity (tree:children node)))))
           (vals (node)
             (cond
               ((constraint-value (d:quantity node))
                 (list (constraint-value (d:quantity node))))
               ((constraint-override (d:quantity node))
                (list (constraint-override (d:quantity node))))
               (t (gen-possibles (d:quantity node)))))
           (ready (arg)
             (or (zerop (length (tree:children arg)))
                 (constraint-value (d:quantity arg)))))
    ;; there are either 1 or 3 arguments
    (cond
      ((eql (length (tree:children node)) 1)
       (dprint `(single-node ,(vals (tree:child node))))
       (loop
          for val in (vals (tree:child node))
          when (valid-val val (list (tree:child node)) (tree:parents node))
          do (with-values (tree:children node)
               (with-values (tree:parents node)
                 (apply-value val (tree:child node))
                 (apply-values val (tree:parents node))
                 (dprint `(single-node-move-on ,val))
                 (move-on)))
            (dprint 'single-try-again))
       (dprint `(single-node-fail ,(vals (tree:child node))))
       nil)
      ((eql (length (tree:children node)) 3)
       ;; boxes with both number arguments having children can only be
       ;; processed when they have a value.
       (if (and (ready (first (tree:children node)))
                (ready (third (tree:children node))))
           (progn
             (dprint `(3-try ,node ,(vals (second (tree:children node)))
                            ,(vals (first (tree:children node)))
                            ,(vals (third (tree:children node)))))
             (loop for i in (vals (second (tree:children node))) do
                  (loop for j in (vals (first (tree:children node))) do
                       (loop
                          for k in (vals (third (tree:children node)))
                          when (and (valid-val (calc i j k) (tree:children node) (tree:parents node))
                                    (valid-box (list j i k) node))
                          do (when (eql k 1/3)
                              (dprint `(gooo ,(length (tree:children (third (tree:children node)))) ,node ,(mapcar #'vals (tree:children node)))))
                            (with-values (tree:children node)
                              (with-values (tree:parents node)
                                (apply-value i (second (tree:children node)))
                                (apply-value j (first (tree:children node)))
                                (apply-value k (third (tree:children node)))
                                (apply-values (funcall i j k) (tree:parents node))
                                (move-on)))
                          (dprint 'they-fucked-up))))
             (dprint `(3-exhausted ,@(mapcar 'd:quantity (tree:children node))))
             nil)
           (return-from walk-template t)))
      (t (error 'generator-error :extra "I'm hardcoded to handle only 1 or 3 arguments.")))))

(defun clean-out-tree (tree keep)
  (check-type tree d:box)
  (assert (eql (length (tree:children tree)) 1))
  (let ((x (d:find-if-on-tree 'xbox-p tree)))
    (if x
        (ecase keep
          (:original)
          (:top
           (d:funcall-on-tree (lambda (node)
                                (when (and (typep node 'd:argument)
                                           (plusp (length (tree:children node)))
                                           (not (eq node (tree:child x))))
                                  (setf (d:quantity node) nil)))
                              tree))
          (:bottom
           (d:funcall-on-tree (lambda (node)
                                (when (and (typep node 'd:argument)
                                           (or (plusp (length (tree:children node)))
                                               (eq node (tree:child x))))
                                  (setf (d:quantity node) nil)))
                              (tree:child (tree:child tree)))))
        (d:funcall-on-tree (lambda (node)
                             (when (and (typep node 'd:argument)
                                        (plusp (length (tree:children node))))
                               (setf (d:quantity node) nil)))
                           (tree:child tree)))
    tree))

(defun fill-in-template (template keep)
  (labels ((copy (template)
             ;; we need a special copy function because the structures aren't explicitely copied
             (prog1 (tree:copy-node template t)
               (d:funcall-on-tree (lambda (a) (when (typep a 'd:argument)
                                                (setf (d:quantity a) (copy-structure (d:quantity a)))))
                                  template)))
           (filled-p (template)
             (not (d:find-if-on-tree (lambda (n)
                                       (and (typep n 'd:argument)
                                            (null (constraint-value (d:quantity n)))))
                                     template)))
           (fill (template)
             (loop
                for sanity-check from 0
                for new = (copy template)
                for leaves = (d:collect-on-tree 'd:leaf-box-p new) do
                (mapc 'walk-template leaves)
                ;; if overrides made it impossible to satisfy the
                ;; constraints, then just toss some values into the
                ;; remaining empty slots.
                (unless (filled-p new)
                  (mapc (lambda (l) (walk-template l t)) leaves))
                (unless (filled-p new)
                  (error 'generator-error
                         :template (d:collect-on-tree (lambda (n) (and (typep n 'd:argument) (null (constraint-value (d:quantity n)))))
                                                      template)
                         :extra "some arguments are still empty."))
                (d:funcall-on-tree (lambda (n)
                                     (when (typep n 'd:argument)
                                       (setf (d:quantity n) (constraint-value (d:quantity n)))))
                                   new)
                until (only-one-answer-p new)
                finally (return new)
                when (>= sanity-check 500) do (error 'generator-error :template template :extra "I tried 500 times but couldn't fill the template in!")))
           (simple-fill (template)
             "just randomly pick arguments based on constraints. Leave numbers untouched."
             (d:funcall-on-tree (lambda (n)
                                  (when (typep n 'd:argument)
                                    (cond
                                      ((typep (d:quantity n) 'op-constraint)
                                       (setf (d:quantity n) (u:random-elt (constraint-ops (d:quantity n)))))
                                      (t
                                       (setf (d:quantity n) (constraint-override (d:quantity n)))))))
                                template)
             template))
    (let ((tree (if (eq keep :original)
                    (simple-fill template)
                    (fill template))))
      (make-problem
       :difficulty (difficulty tree)
       :tree (clean-out-tree tree keep)))))

(defun only-one-answer-p (tree)
  "If the X box at the top is multiplied by two values in two
different paths and in the box that joins the two paths their sum is
zero then any number can satisfy X. This returns T if there is only
one answer for X and NIL otherwise."
  (labels ((sub (box value)
             (dolist (p (tree:parents box))
               (check-type p d:argument)
               (setf (d:quantity p) value)
               (sub (tree:parent p)
                    (let ((c (tree:children (tree:parent p))))
                      (ecase (length c)
                        (1 (d:quantity p))
                        (3 (funcall (d:quantity (second c))
                                    (d:quantity (first c))
                                    (d:quantity (third c))))))))))
    (setf tree (tree:copy-node tree t))
    (let* ((bottom (d:quantity (tree:child tree)))
           (x (d:find-if-on-tree 'xbox-p tree))
           (top (and x (d:quantity (tree:child x)))))
      ;; if there's no X then we have nothing to worry about.
      (or (not x)
          ;; this works because there are no trees where X is multiplied
          ;; by itself.
          (loop for i in (loop for i from 1 to 10 unless (eql i top) collect i)
             do
             (handler-case (sub x i)
               ;; we can safely ignore this because the bottom won't be
               ;; changed so it'll inherit the last value, which was
               ;; valid.
               (division-by-zero ()))
             (when (eql bottom (d:quantity (tree:child tree))) (dprint `(caught-one ,top ,i ,bottom ,(d:quantity (tree:child tree)) ,tree)))
             never (eql bottom (d:quantity (tree:child tree))))))))

(defun expand-template (template)
  (let ((x-box [m [a (num)]]))
    (labels ((expand (n)
               (cond
                 ((string= n 'n) (num))
                 ((or (string= n 'o) (string= n 'oo)) (op))
                 ((or (string= n 'a) (string= n 'aa)) (make-op-constraint :ops '(+ -)))
                 ((or (string= n 'm) (string= n 'mm)) (make-op-constraint :ops '(* /)))
                 ((or (string= n '*) (string= n '**)) (make-op-constraint :ops '(*)))
                 ((or (string= n '/) (string= n '//)) (make-op-constraint :ops '(/)))
                 ((or (string= n '+) (string= n '++)) (make-op-constraint :ops '(+)))
                 ((or (string= n '-) (string= n '--)) (make-op-constraint :ops '(-)))
                 (t (error 'generator-error :template template :extra (format nil "Unknown symbol ~s" n)))))
             (recurse (template)
               (cond
                 ((listp template)
                  (let ((first (recurse (first template)))
                        (second (recurse (second template)))
                        (third (recurse (third template)))
                        (swap (eql (length (symbol-name (second template))) 2)))
                    (loop for i in first nconc
                         (loop for j in second nconc
                              (loop for k in third
                                 collect [a (num) [m i j k]]
                                 if swap
                                 collect [a (num) [m k j i]])))))
                 ((and (symbolp template)
                       (string= template 'x))
                  (list [a (num) x-box]))
                 ((symbolp template)
                  (list [a (expand template)]))
                 ((numberp template)
                  (list [a (make-num-constraint :override template)]))
                 (t (error 'generator-error :template template :extra "dunno how to handle this.")))))
      (let ((trees (recurse template)))
        (loop for i in trees
           collect [m i])))))

(defun expand-templates (&rest templates)
  (mapcan 'expand-template templates))

(defun expand-and-fill (keep &rest templates)
  (dprint `(====== ,(get-internal-real-time) ======))
  (fill-in-template
   ;; don't let templates with more permutations skew the
   ;; probabability of it being picked.
   (u:random-elt
    (u:random-elt
     (mapcar 'expand-template templates)))
   keep))

(defun difficulty (tree)
  "Return a number representing the difficulty. The higher the number, the harder the problem."
  (print `(diff ,tree))
  (labels ((count (fn node)
             (length (d:collect-on-tree (lambda (n)
                                          (and (typep n 'd:argument)
                                               (d:quantity n)
                                               (funcall fn (d:quantity n))))
                                        node))))
    (+ (* 1 (count (lambda (q) (and (typep q 'number) (minusp q))) tree))
       (* 3 (count (lambda (q) (typep q 'ratio)) tree))
       (* 2 (count (lambda (q) (eq q '/)) tree))
       (* 2 (count (lambda (q) (eq q '-)) tree))
       (* 4 (length (d:collect-on-tree (constantly t) tree :only-boxes t))))))

(defun expand-and-fill-n-times (n keep &rest templates)
  (mapcar 'problem-tree
          (sort
           (loop
              while (< (length acc) n)
              for p = (apply 'expand-and-fill keep templates)
              for sanity-check from 0
              unless (find (problem-tree p) acc :test 'tree:node-equal :key 'problem-tree)
              collect p into acc
              when (>= sanity-check 500) do (error 'generator-error :template templates :extra (format nil "I tried 500 times and only produced ~d of ~d expansions." (length acc) n))
              finally (return acc))
           '<
           :key 'problem-difficulty)))
